package com.kunge2013.chatbi.chat;

import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.zhipuai.ZhiPuAiChatModel;
import org.springframework.ai.zhipuai.ZhiPuAiChatOptions;
import org.springframework.ai.zhipuai.api.ZhiPuAiApi;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.publisher.Flux;

import java.time.Duration;
import java.time.LocalDateTime;
import java.util.Map;

/**
 * @Author kunge2013
 * @Description 聊天模型
 * @Date 2024/11/11 15:32
 * @Version 1.0
 */
@RestController
@RequestMapping("/chatApi")
@AllArgsConstructor
@Slf4j
public class ChatController {

    private final ZhiPuAiChatModel chatModel;

    @PostMapping(path = "chat/{message}")
    public Map chat(@PathVariable(name = "message") String message) {
        return Map.of("generation", this.chatModel.call(message));
    }

    @GetMapping("/ai/generateStream")
    public Flux<ChatResponse> generateStream(@RequestParam(value = "message", defaultValue = "给我考个玩笑") String message) {
        var prompt = new Prompt(new UserMessage(message));
        return this.chatModel.stream(prompt);
    }

    private static final Long CHAT_TIMEOUT = Duration.ofMinutes(50).toMillis();

    @GetMapping("/custom/stream/chat")
    @CrossOrigin
    public SseEmitter customChat() {
        SseEmitter emitter = new SseEmitter(CHAT_TIMEOUT);

        // Set event handler for SSEEmitter
        emitter.onCompletion(() -> log.info(LocalDateTime.now() + ", on completion"));
        emitter.onTimeout(() -> {
            log.info(LocalDateTime.now() + ", uid# on timeout");
            emitter.complete();
        });

        // Start a new thread to generate SSE events
        new Thread(() -> {
            try {
                var prompt = new Prompt(new UserMessage("给我考个玩笑"));
                Flux<ChatResponse> stream = this.chatModel.stream(prompt);
                for (int i = 0; i < 10; i++) {
                    emitter.send(SseEmitter.event().name("message").data("Event " + i));
                    Thread.sleep(1000);
                }
            } catch (Exception e) {
                emitter.completeWithError(e);
            } finally {
                emitter.complete();
            }
        }).start();
        return emitter;
    }

}
